package org.yray.tts.dic;

import lombok.extern.slf4j.Slf4j;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.nio.charset.Charset;
import java.util.HashMap;
import java.util.Locale;
import java.util.Map;

@Slf4j
public class WordDict {
    private DictSegment dictSegment= new DictSegment((char) 0);
    private Double total = 0.0;
    public final Map<String, Double> freqs = new HashMap<>();
    private Double minFreq = Double.MAX_VALUE;
    public static final  double DEF_FREQ=3.0;
    private static final String MAIN_DICT = "/dict.txt";
    private static String USER_DICT_SUFFIX = ".dict";
    private String addWordToSegment(String word) {
        if (null != word && !"".equals(word.trim())) {
            String key = word.trim();
            dictSegment.fillSegment(key.toCharArray());
            return key;
        }
        else {
            return null;
        }
    }

    public void addWord(String word,double freq)
    {
        total += freq;
        word = addWordToSegment(word);
        freqs.put(word, freq);
    }

    public void addWord(String word)
    {
        addWord(word,DEF_FREQ);
    }

    //归一化
    public void normalize()
    {
        // normalize
        for (Map.Entry<String, Double> entry : freqs.entrySet()) {
            entry.setValue((Math.log(entry.getValue() / total)));
            minFreq = Math.min(entry.getValue(), minFreq);
        }
    }

    public DictSegment getTrie() {
        return this.dictSegment;
    }

    public boolean containsWord(String word) {
        return freqs.containsKey(word);
    }


    public Double getFreq(String key) {
        if (containsWord(key)) {
            return freqs.get(key);
        }
        else {
            return minFreq;
        }
    }

    public void clean()
    {
        freqs.clear();
    }
    public DictSegment getDictSegment() {
        return dictSegment;
    }


    public void loadDict() {
        dictSegment = new DictSegment((char) 0);
        InputStream is = this.getClass().getResourceAsStream(MAIN_DICT);
        try {
            BufferedReader br = new BufferedReader(new InputStreamReader(is, Charset.forName("UTF-8")));

            long s = System.currentTimeMillis();
            while (br.ready()) {
                String line = br.readLine();
                String[] tokens = line.split("[\t ]+");

                if (tokens.length < 2) {
                    continue;
                }

                String word = tokens[0];
                double freq = Double.valueOf(tokens[1]);
                 addWord(word,freq);
            }
            // normalize
            for (Map.Entry<String, Double> entry : freqs.entrySet()) {
                entry.setValue((Math.log(entry.getValue() / total)));
                minFreq = Math.min(entry.getValue(), minFreq);
            }
            log.info(String.format("main dict load finished, time elapsed {} ms",
                    System.currentTimeMillis() - s));
        }
        catch (IOException e) {
            log.error("{} load failure!", MAIN_DICT);
        }
        finally {
            try {
                if (null != is) {
                    is.close();
                }
            }
            catch (IOException e) {
                log.error("{} close failure!", MAIN_DICT);
            }
        }
    }
}
